from langchain.agents import StructuredChatAgent, AgentExecutor
from langchain.memory import ConversationBufferMemory
from langchain.tools import tool
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI


# --------------------------
# 1. 工具定义（修复审批人识别逻辑+流程顺序）
# --------------------------
@tool(description="查询部门预算是否充足。参数: department(部门名), amount(金额)。返回：预算充足/不足，当前剩余预算")
def check_budget_tool(department: str, amount: float) -> str:
    budget_data = {
        "技术部": 10000,
        "市场部": 8000,
        "行政部": 5000
    }
    remaining = budget_data.get(department, 0)
    if remaining >= amount:
        return f"预算充足。{department}当前剩余预算：{remaining}元，申请金额：{amount}元"
    else:
        return f"预算不足。{department}当前剩余预算：{remaining}元，申请金额：{amount}元"


@tool(description="提交审批。参数: approver(审批人，如技术部经理、CEO), application(申请内容)。返回：审批结果（通过/驳回）")
def submit_approval_tool(approver: str, application: str) -> str:
    """修复：灵活识别审批人（支持部门+经理/主管，如技术部经理）"""
    # 新增：判断“部门经理/主管”（不管前面带哪个部门）
    is_department_manager = any(keyword in approver for keyword in ["经理", "主管"])
    is_ceo = "CEO" in approver or "首席执行官" in approver

    if is_department_manager:
        return f"{approver}审批通过：{application}"
    elif is_ceo:
        return f"{approver}审批通过：{application}"
    else:
        return f"{approver}审批驳回：未知审批人（仅支持“XX经理/主管”或“CEO”）"


@tool(description="财务打款。参数: recipient(收款人), amount(金额)。返回：打款结果")
def finance_payment_tool(recipient: str, amount: float) -> str:
    return f"财务已向{recipient}打款{amount}元，交易单号：PAY{hash(recipient + str(amount))}"


@tool(
    description="查询当前流程状态。参数: application_id(申请ID)。注意：仅在用户主动要求查状态时调用，提交申请后第一步禁止调用")
def query_status_tool(application_id: str) -> str:
    from langchain_core.messages import get_buffer_string  # 正确导入路径
    history = memory.load_memory_variables({})["chat_history"]
    history_str = get_buffer_string(history)
    # 只显示最近300字符，避免信息过长
    return f"申请ID {application_id} 的流程状态：\n{history_str[-300:]}"


tools = [check_budget_tool, submit_approval_tool, finance_payment_tool, query_status_tool]

# --------------------------
# 2. LLM和记忆初始化（不变）
# --------------------------
llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key="sk-6434b554122c4b049ceb805d703f695b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True
)

# --------------------------
# 3. 提示词模板（修复：强制流程顺序）
# --------------------------
prompt_template = """
你是采购申请审批工作流的处理Agent，必须严格遵守以下流程顺序，禁止跳步：

1. 【强制流程步骤】
   - 第一步（必走）：收到采购申请后，无论任何情况，必须先调用check_budget_tool检查部门预算；
   - 第二步（预算充足才走）：
     - 金额≤5000元 → 调用submit_approval_tool提交“XX部门经理/主管”审批；
     - 金额>5000元 → 先提交“XX部门经理/主管”审批，再提交“CEO”审批；
   - 第三步（所有审批通过才走）：调用finance_payment_tool执行打款；
   - 【禁止】提交申请后第一步不能调用query_status_tool（仅用户主动查状态时用）。

2. 工具调用格式：{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}
3. 对话历史：{chat_history}

用户的采购申请：{input}

{agent_scratchpad}
"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["input", "chat_history", "agent_scratchpad"]
)

# --------------------------
# 4. 创建Agent并运行
# --------------------------
agent = StructuredChatAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True
)

agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    memory=memory,
    verbose=True,
    handle_parsing_errors=True
)

# --------------------------
# 测试
# --------------------------
if __name__ == "__main__":
    print("===== 提交采购申请 =====")
    res1 = agent_executor.invoke({
        "input": "申请ID: PUR2023001，技术部采购服务器，金额6000元，原因：业务扩容"
    })
    print("结果：", res1["output"], "\n")

    print("===== 查询流程状态 =====")
    res2 = agent_executor.invoke({
        "input": "查询申请ID: PUR2023001的当前进度"
    })
    print("结果：", res2["output"], "\n")